import argparse
import os
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from models.source_fcanet import fcanet34, fcanet50, fcanet101, fcanet152

# 解析命令行参数
parser = argparse.ArgumentParser(description='Train FcaNet model')
parser.add_argument('--model', type=str, required=True, choices=['fcanet34', 'fcanet50', 'fcanet101', 'fcanet152'], help='Specify the model to use: fcanet34, fcanet50, fcanet101, fcanet152')
parser.add_argument('--strategy', type=str, required=True, choices=['top1', 'top2', 'top4', 'top8', 'top16', 'top32', 'bot1', 'bot2', 'bot4', 'bot8', 'bot16', 'bot32', 'low1', 'low2', 'low4', 'low8', 'low16', 'low32'], help='Specify the frequency selection strategy')
args = parser.parse_args()

# 定义模型保存路径
save_path = './results'
if not os.path.exists(save_path):
    os.makedirs(save_path)
def load_data():
    transform = transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])

    trainset = torchvision.datasets.CIFAR100(root='../data', train=True, download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True, num_workers=4)

    testset = torchvision.datasets.CIFAR100(root='../data', train=False, download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(testset, batch_size=32, shuffle=False, num_workers=4)

    return trainloader, testloader

# 定义设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 根据输入参数选择模型
if args.model == 'fcanet34':
    model = fcanet34(num_classes=100, strategy=args.strategy)
elif args.model == 'fcanet50':
    model = fcanet50(num_classes=100, strategy=args.strategy)
elif args.model == 'fcanet101':
    model = fcanet101(num_classes=100, strategy=args.strategy)
elif args.model == 'fcanet152':
    model = fcanet152(num_classes=100, strategy=args.strategy)

model = model.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

# 训练函数
def train(epoch,trainloader):
    model.train()
    running_loss = 0.0
    for batch_idx, (inputs, targets) in enumerate(trainloader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if batch_idx % 100 == 99:
            print(f'Epoch [{epoch + 1}], Batch [{batch_idx + 1}], Loss: {running_loss / 100:.4f}')
            running_loss = 0.0

# 测试函数
def test(testloader):
    model.eval()
    correct = 0
    total = 0
    test_loss = 0.0
    with torch.no_grad():
        for inputs, targets in testloader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    accuracy = 100.0 * correct / total
    test_loss /= len(testloader)
    print(f'Test Accuracy: {accuracy:.2f}%')
    return accuracy, test_loss


# 训练循环
if __name__ == "__main__":
    num_epochs = 100
    trainloader, testloader = load_data()
    best_accuracy = 0.0
    accuracies = []
    best_model_path = './results/best_model-Fcanet.pth'
    for epoch in range(num_epochs):
        train(epoch,trainloader)
        accuracy, test_loss=test(testloader)
        scheduler.step()
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            torch.save(model.state_dict(), best_model_path)
            print(f'Saved the best model with accuracy: {best_accuracy:.2f}%')
        accuracies.append(accuracy)

    print('Finished Training')
